/***************************************************************************
 *
 *   Copyright (C) 2025 by Will Gauvin
 *   Licensed under the Academic Free License version 2.1
 *
 ***************************************************************************/

#include "dsp/RescaleMeanStdCalculator.h"

#include <assert.h>
#include <algorithm>

using namespace std;

template<typename T>
void zero (vector<T>& data)
{
  std::fill(data.begin(), data.end(), 0);
}

void dsp::RescaleMeanStdCalculator::init(const dsp::TimeSeries* input, uint64_t nsample, bool output_time_total)
{
  npol  = input->get_npol();
  nchan = input->get_nchan();
  ndim = input->get_ndim();

  if (output_time_total)
    time_total.resize (npol);

  freq_total.resize (npol);
  freq_totalsq.resize (npol);

  scale.resize (npol);
  offset.resize (npol);

  for (unsigned ipol=0; ipol < npol; ipol++)
  {
    if (output_time_total)
    {
      time_total[ipol].resize (nsample);
      zero (time_total[ipol]);
    }

    freq_total[ipol].resize (nchan);
    zero (freq_total[ipol]);

    freq_totalsq[ipol].resize (nchan);
    zero (freq_total[ipol]);

    scale[ipol].resize (nchan);
    offset[ipol].resize (nchan);
  }
}

uint64_t dsp::RescaleMeanStdCalculator::sample_data(const dsp::TimeSeries* input, uint64_t start_dat, uint64_t end_dat, bool output_time_total)
{
  const bool input_voltage = (input->get_state() == Signal::Nyquist || input->get_state() == Signal::Analytic);
  uint64_t samp_dat = start_dat;

  switch (input->get_order())
  {
  case TimeSeries::OrderTFP:
  {
    const float* in_data = input->get_dattfp();
    in_data += start_dat * nchan * npol * ndim;
    for (uint64_t idat=start_dat; idat < end_dat; idat++, samp_dat++)
    {
      for (unsigned ichan=0; ichan < nchan; ichan++)
      {
        for (unsigned ipol=0; ipol < npol; ipol++)
        {
          for (unsigned idim=0; idim < ndim; idim++)
          {
            const double val = static_cast<double>(*in_data);
            freq_total[ipol][ichan] += val;
            freq_totalsq[ipol][ichan] += val * val;

            if (output_time_total)
            {
              if (input_voltage)
                time_total[ipol][samp_dat] += static_cast<float>(val * val);
              else
                time_total[ipol][samp_dat] += static_cast<float>(val);
            }
            in_data++;
          }
        }
      }
    }
    break;
  }
  case TimeSeries::OrderFPT:
  {
    for (unsigned ipol=0; ipol < npol; ipol++)
    {
      for (unsigned ichan=0; ichan < nchan; ichan++)
      {
        const float* in_data = input->get_datptr (ichan, ipol);

        samp_dat = start_dat;

        double sum = 0.0;
        double sumsq = 0.0;

        for (uint64_t idat=start_dat; idat < end_dat; idat++, samp_dat++)
        {
          for (unsigned idim=0; idim<ndim; idim++)
          {
            const double val = double(*in_data);
            sum += val;
            sumsq += val * val;

            if (output_time_total)
            {
              if (input_voltage)
                time_total[ipol][samp_dat] += static_cast<float>(val * val);
              else
                time_total[ipol][samp_dat] += static_cast<float>(val);
            }
            in_data++;
          }
        }
        freq_total[ipol][ichan] += sum;
        freq_totalsq[ipol][ichan] += sumsq;
      }
    }
    break;
  }
  default:
    throw Error (InvalidState, "dsp::Rescale::operate",
        "Requires data in TFP or FPT order");
  }

  return samp_dat;
}

void dsp::RescaleMeanStdCalculator::compute(uint64_t nsample)
{
  if (nsample == 0)
    throw Error (InvalidParam, "dsp::RescaleMeanStdCalculator::compute",
		 "invalid nsample=" I64, nsample);

  const double nval = static_cast<double>(nsample * ndim);

  for (unsigned ipol=0; ipol < npol; ipol++)
  {
    for (unsigned ichan=0; ichan < nchan; ichan++)
    {
      double mean = freq_total[ipol][ichan] / nval;
      double meansq = freq_totalsq[ipol][ichan] / nval;
      double variance = meansq - mean*mean;

      offset[ipol][ichan] = -mean;
      if (variance <= 0.0)
        scale[ipol][ichan] = 1.0;
      else
        scale[ipol][ichan] = 1.0 / sqrt(variance);
    }
  }
}

void dsp::RescaleMeanStdCalculator::reset_sample_data()
{
  for (unsigned ipol=0; ipol < npol; ipol++)
  {
    zero (freq_total[ipol]);
    zero (freq_totalsq[ipol]);
    if (ipol < time_total.size())
      zero (time_total[ipol]);
  }
}

const double* dsp::RescaleMeanStdCalculator::get_mean (unsigned ipol) const
{
  assert (ipol < freq_total.size());
  return &(freq_total[ipol][0]);
}

const double* dsp::RescaleMeanStdCalculator::get_variance (unsigned ipol) const
{
  assert (ipol < freq_totalsq.size());
  return &(freq_totalsq[ipol][0]);
}
